import os, sys, logging
sys.path.append(".")
import matplotlib.pyplot as plt

import pyarma as pa
import numpy as np
import time as tm

from unifuncs import *
from expoint import *
from problems import *

def Ureal(M, lb, ub, t):

    h = (ub - lb) / M

    u = pa.mat(M, 1)
    u[0] = t
    for i in range(1, M):
        x = lb + i * h
        u[i] = x * (1 - x) * np.exp(2 * t)

    return u

def HeatEq_F(u0, M, lb, ub):

    h = (ub - lb) / M
    un = pa.zeros(pa.size(u0))
    un[0] = 1
    for i in range(1, M):
        x = lb + i * h
        un[i] = 2 * np.exp(2 * u0[0])

    return un

def HeatEq_J(u0, M, lb, ub):

    h = (ub - lb) / M
    Jn = pa.zeros(M, M)
    for i in range(1, M):
        x = lb + i * h
        Jn[i, 0] = 4 * np.exp(2 * u0[0])
    
    Jn[1:M-1, 1:M-1] = diff.cd2(M - 1, h) + 2 * pa.eye(M - 1, M - 1)
    return Jn

if __name__ == '__main__':

    logger = logging.getLogger(__name__)
    logger.setLevel(level = logging.INFO)
    handler = logging.FileHandler("./logs/logs_heateq.txt")
    handler.setLevel(level = logging.INFO)
    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    handler.setFormatter(formatter)
    logger.addHandler(handler)

    M = 16
    lb = 0.0
    ub = 1.0
    h = (ub - lb) / M

    A = pa.zeros(M, M)
    A[1:M-1, 1:M-1] = diff.cd2(M - 1, h) + 2 * pa.eye(M - 1, M - 1)
    
    u_real = lambda t: Ureal(M, lb, ub, t)

    g = lambda u: HeatEq_F(u, M, lb, ub)
    F = lambda u: A * u + g(u)
    J = lambda u: HeatEq_J(u, M, lb, ub)

    Ns = [(2 ** (8+i)) for i in range(4)]
    RKs = [RK.Euler1, RK.StrehmelStrong2, RK.StrehmelWeak2, RK.Cox4, RK.Krogstad4, RK.Strehmel4, RK.Hochbruck5]
    RBs = [RB.Euler1, RB.StrehmelStrong2, RB.StrehmelWeak2, RB.Cox4, RB.Krogstad4, RB.Strehmel4, RB.Hochbruck5]
    
    logger.info("Simulation started. TimeMesh = {}".format(Ns))

    for i in range(len(RKs)):
        
        fig1, ax1 = plt.subplots()
        ax1.set(xlabel=r'$\log \tau$', ylabel = r'$\log L_2$')
        fig2, ax2 = plt.subplots()
        ax2.set(xlabel=r'$\log \tau$', ylabel = r'$\log L_\infty$')
        
        euk2 = []
        euki = []
        erk2 = []
        erb2 = []
        erki = []
        erbi = []
        taus = []
        for N in Ns:
            tau = 1 / N
            taus.append(tau)
            u0uk = u_real(0)
            flag_uk = True
            u0rk = u_real(0)
            flag_rk = True
            u0rb = u_real(0)
            flag_rb = True
            for j in range(N):
                if flag_uk == True:
                    u0uk = RKs[i](u0uk, tau, pa.zeros(M, M), F)
                    if np.sqrt(h) * pa.norm(pa.abs(u0uk - u_real((j + 1) * tau)), 2) > 1 or pa.norm(pa.abs(u0uk - u_real((j + 1) * tau)), 'inf') > 1:
                        euk2.append("不收敛")
                        euki.append("不收敛")
                        flag_uk = False
                if flag_rk == True:
                    u0rk = RKs[i](u0rk, tau, A, g)
                    if np.sqrt(h) * pa.norm(pa.abs(u0rk - u_real((j + 1) * tau)), 2) > 1 or pa.norm(pa.abs(u0rk - u_real((j + 1) * tau)), 'inf') > 1:
                        erk2.append("不收敛")
                        erki.append("不收敛")
                        flag_rk = False
                if flag_rb == True:
                    u0rb = RBs[i](u0rb, tau, J, F)
                    if np.sqrt(h) * pa.norm(pa.abs(u0rb - u_real((j + 1) * tau)), 2) > 1 or pa.norm(pa.abs(u0rb - u_real((j + 1) * tau)), 'inf') > 1:
                        erb2.append("不收敛")
                        erbi.append("不收敛")
                        flag_rk = False
            if flag_uk == True:
                euk2.append(np.sqrt(h) * pa.norm(pa.abs(u0uk - u_real(1)), 2))
                euki.append(pa.norm(pa.abs(u0uk - u_real(1)), 'inf'))
            if flag_rk == True:
                erk2.append(np.sqrt(h) * pa.norm(pa.abs(u0rk - u_real(1)), 2))
                erki.append(pa.norm(pa.abs(u0rk - u_real(1)), 'inf'))
            if flag_rb == True:
                erb2.append(np.sqrt(h) * pa.norm(pa.abs(u0rb - u_real(1)), 2))
                erbi.append(pa.norm(pa.abs(u0rb - u_real(1)), 'inf'))
                
        logurk = "Underlying\tRungeKutta\t" + RKs[i].__name__ + ":\nL2 = "
        logerk = "Exponential\tRungeKutta\t" + RKs[i].__name__ + ":\nL2 = "
        logerb = "Exponential\tRosenbrock\t" + RBs[i].__name__ + ":\nL2 = "
        
        x2uk, y2uk = [], []
        x2rk, y2rk = [], []
        x2rb, y2rb = [], []
        
        for i in range(len(Ns)):
            
            if euk2[i] == '不收敛':
                logurk += "%s" % euk2[i]
            else:
                logurk += "%.4e" % euk2[i]
                x2uk.append(np.log(taus[i]))
                y2uk.append(np.log(euk2[i]))
            if i < len(Ns) - 1:
                logurk += "\t&\t"
                
            if erk2[i] == '不收敛':
                logerk += "%s" % erk2[i]
            else:
                logerk += "%.4e" % erk2[i]
                x2rk.append(np.log(taus[i]))
                y2rk.append(np.log(erk2[i]))
            if i < len(Ns) - 1:
                logerk += "\t&\t"
                
            if erb2[i] == '不收敛':
                logerb += "%s" % erb2[i]
            else:
                logerb += "%.4e" % erb2[i]
                x2rb.append(np.log(taus[i]))
                y2rb.append(np.log(erb2[i]))
            if i < len(Ns) - 1:
                logerb += "\t&\t"
                
        ouk, buk = np.polyfit(x2uk, y2uk, 1)
        ork, brk = np.polyfit(x2rk, y2rk, 1)
        orb, brb = np.polyfit(x2rb, y2rb, 1)
                
        logurk += "\t&\t%.4lf\nLinf = " % ouk
        logerk += "\t&\t%.4lf\nLinf = " % ork
        logerb += "\t&\t%.4lf\nLinf = " % orb
        
        xiuk, yiuk = [], []
        xirk, yirk = [], []
        xirb, yirb = [], []
        
        for i in range(len(Ns)):
            
            if euki[i] == '不收敛':
                logurk += "%s" % euki[i]
            else:
                logurk += "%.4e" % euki[i]
                xiuk.append(np.log(taus[i]))
                yiuk.append(np.log(euki[i]))
            if i < len(Ns) - 1:
                logurk += "\t&\t"
                
            if erki[i] == '不收敛':
                logerk += "%s" % erki[i]
            else:
                logerk += "%.4e" % erki[i]
                xirk.append(np.log(taus[i]))
                yirk.append(np.log(erki[i]))
            if i < len(Ns) - 1:
                logerk += "\t&\t"
                
            if erbi[i] == '不收敛':
                logerb += "%s" % erbi[i]
            else:
                logerb += "%.4e" % erbi[i]
                xirb.append(np.log(taus[i]))
                yirb.append(np.log(erbi[i]))
            if i < len(Ns) - 1:
                logerb += "\t&\t"
        
        ouk, buk = np.polyfit(xiuk, yiuk, 1)
        ork, brk = np.polyfit(xirk, yirk, 1)
        orb, brb = np.polyfit(xirb, yirb, 1)
                
        logger.info(logurk + "\t&\t%.4lf" % ouk)
        logger.info(logerk + "\t&\t%.4lf" % ork)
        logger.info(logerb + "\t&\t%.4lf" % orb)
        
        ax1.plot(x2uk, y2uk, "^-", label="Underlying Runge-Kutta Method")
        ax1.plot(x2rk, y2rk, "s-", label="Exponential Runge-Kutta Method")
        ax1.plot(x2rb, y2rb, "*-", label="Exponential Rosenbrock Method")
        ax1.legend()
        
        ax2.plot(xiuk, yiuk, "^-", label="Underlying Runge-Kutta Method")
        ax2.plot(xirk, yirk, "s-", label="Exponential Runge-Kutta Method")
        ax2.plot(xirb, yirb, "*-", label="Exponential Rosenbrock Method")
        ax2.legend()
        
        plt.show()
        
    logger.info("Simulation ended.\n")